#!/usr/bin/env python

import cPickle
import glob
import optparse
import os
import sys

import histogram

import batchcat
import charmicat

def main():
    """
    bcat_module_hists.py [opts] file1.h5 [file2.h5 [...]]

    Fill histograms via batch submission from a module supplied from the command 
    line or inferred from the file.
    """

    op = optparse.OptionParser(usage=main.__doc__)

    op.add_option("-f", "--force", dest="force", default=False, action='store_true',
                  help="Force overwrite of output directory")
    op.add_option("-l", "--logdir", dest='logdir', default=None,
                  help='Specify directory PATH for MP logging', metavar='PATH')
    op.add_option("-m", "--module", dest='module_name', default=None, 
                  help="Use module MODULE instead of inferring from input", metavar='MODULE')
    op.add_option("-n", "--name", dest="list", action='append', default=[],
                  help="Fill histogram NAME (can be used multiple times, compatible w/ --histset",
                  metavar="NAME")
    op.add_option("-o", "--output", dest="outdir", default=None, 
                  help="Leave output in DIR", metavar="DIR")
    op.add_option("-v", "--verbose", dest='verbose', default=False, action='store_true',
                  help='Print more output')
    op.add_option("--cols-only", dest='columns_only', action='store_true', default=False,
                  help="Only fill histograms whose associated columns are present in the input tables")
    op.add_option("--njobs", dest="njobs", type="int", default=100,
                  help="Spawn no more than N threads", metavar="N")
    op.add_option("--force-log", dest='force_log', default=False, action='store_true',
                  help='Force reuse of currently-existing log directory')
    op.add_option("--histset", dest='histset', default=None,
                  help="Specify a subset SET of plots to fill", metavar="SET")

    batchcat.add_options(op)

    opts, args = op.parse_args()

    if len(args) == 0:
        msg = "Please supply at least one input file."
        raise charmicat.BadCommandLineError(msg, -1)

    if not opts.module_name:
        opts.module_name = charmicat.module_name_from_paths(args)
    
    try:
        module = charmicat.safe_import(opts.module_name)           
    except ImportError, e:
        msg = "Error importing module `{0}': \n{1}".format(opts.module_name, e)
        raise charmicat.BadCommandLineError(msg, -2)

    hnames = []

    if not opts.histset and len(opts.list) == 0:
        opts.histset = 'default'

    if opts.histset:
        try:
            hnames = module.histsets[opts.histset]
        except AttributeError, e:
            msg = "Error getting histograms from module `{0}':\n{1}".format(opts.module_name, e)
            raise charmicat.BadCommandLineError(msg, -3)
    
    if len(opts.list) != 0:
        hnames = opts.list
        
    for hn in hnames:
        if hn not in module.hists:
            msg = "Histogram `{0}' of histset `{1}' not found in `{1}.hists'".format(hn, opts.module_name)
            raise charmicat.BadCommandLineError(msg, -4)
    
    if opts.outdir is None:
        opts.outdir = os.path.basename(os.getcwd())

    if len(args) < opts.njobs:
        opts.njobs = len(args)
    
    opts.outdir = os.path.abspath(opts.outdir)

    table_columns = charmicat.column_names_from_paths(args)
    hcol_names, hvec_names = charmicat.how_fill(module, hnames, table_columns)
    
    charmicat.assert_fillable(module, hvec_names, table_columns)

    opts.outdir = charmicat.safe_mkdir(opts.outdir, force=opts.force)
    

    hnames_path = os.path.join(opts.outdir, 'hnames.pkl')
    with open(hnames_path, 'w') as fph:
        cPickle.dump((hcol_names, hvec_names), fph)

    if opts.verbose:
        print "Column-fillable histograms: "
        charmicat.column_output(hcol_names, prefix='\t')
        print "Vector-filled histograms: "
        charmicat.column_output(hvec_names, prefix='\t')
     
    rm = batchcat.resource_managers['default']
    
    for i, pathtup in enumerate(charmicat.n_groups(opts.njobs, args)):
        pklpath = os.path.join(opts.outdir, 'job_{0}_paths.pkl'.format(i))
        with open(pklpath, 'w') as fp:
            cPickle.dump(pathtup, fp)

        command_line = 'python -m batchcat.run.module_hists {0} {1} {2}'
        command_line = command_line.format(opts.module_name, 
                                           os.path.abspath(pklpath), 
                                           os.path.abspath(hnames_path))
                                
        rm.add_job('bcat_module_hists' + str(i),
                   command_line, 
                   opts)
        
    if opts.bcat_write is None:
        rm.submit()
    else:
        rm.write(path=opts.bcat_write)



if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)

